// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/testing/exec/util/assert_query_builder.h>
#include <pollux/testing/exec/util/operator_test_base.h>
#include <pollux/plan/plan_builder.h>

using namespace kumo::pollux;
using namespace kumo::pollux::test;
using namespace kumo::pollux::exec::test;

class MarkDistinctTest : public OperatorTestBase {
public:
    void runBasicTest(const VectorPtr &base) {
        const vector_size_t size = base->size() * 2;
        auto indices = make_indices(size, [](auto row) { return row / 2; });
        auto data = wrap_in_dictionary(indices, size, base);

        auto isDistinct =
                make_flat_vector<bool>(size, [&](auto row) { return row % 2 == 0; });

        auto expectedResults = make_row_vector({data, isDistinct});

        auto plan = PlanBuilder()
                .values({make_row_vector({data})})
                .markDistinct("c0_distinct", {"c0"})
                .planNode();

        auto results = AssertQueryBuilder(plan).copyResults(pool());
        assertEqualVectors(expectedResults, results);
    }
};

template<typename T>
class MarkDistinctPODTest : public MarkDistinctTest {
};

using MyTypes = ::testing::Types<int16_t, int32_t, int64_t, float, double>;
TYPED_TEST_SUITE(MarkDistinctPODTest, MyTypes);

TYPED_TEST(MarkDistinctPODTest, basic) {
    auto data = VectorBuilder::make_flat_vector<TypeParam>(
        1'000, [](auto row) { return row; }, [](auto row) { return row == 7; });

    MarkDistinctTest::runBasicTest(data);
}

TEST_F(MarkDistinctTest, tinyint) {
    auto data = make_flat_vector<int8_t>(
        256, [](auto row) { return row; }, [](auto row) { return row == 7; });

    runBasicTest(data);
}

TEST_F(MarkDistinctTest, boolean) {
    auto data = make_nullable_flat_vector<bool>({true, false, std::nullopt});

    runBasicTest(data);
}

TEST_F(MarkDistinctTest, varchar) {
    auto base = make_flat_vector<StringView>({
        "{1, 2, 3, 4, 5}",
        "{1, 2, 3}",
    });
    runBasicTest(base);
}

TEST_F(MarkDistinctTest, array) {
    auto base = make_array_vector<int64_t>({
        {1, 2, 3, 4, 5},
        {1, 2, 3},
    });
    runBasicTest(base);
}

TEST_F(MarkDistinctTest, map) {
    auto base = make_map_vector<int8_t, int32_t>(
        {{{1, 1}, {1, 1}, {1, 1}, {1, 1}, {1, 1}}, {{1, 1}, {1, 1}, {1, 1}}});
    runBasicTest(base);
}

TEST_F(MarkDistinctTest, row) {
    auto base = make_row_vector({
        make_array_vector<int64_t>({
            {1, 2, 3, 4, 5},
            {1, 2, 3},
        }),
        make_map_vector<int8_t, int32_t>({
            {{1, 1}, {1, 1}, {1, 1}, {1, 1}, {1, 1}},
            {{1, 1}, {1, 1}, {1, 1}},
        }),
    });
    runBasicTest(base);
}

TEST_F(MarkDistinctTest, aggregation) {
    // Simulate the input over 3 splits.
    std::vector<RowVectorPtr> vectors = {
        make_row_vector({
            make_flat_vector<int32_t>({1, 1}),
            make_flat_vector<int32_t>({1, 1}),
            make_flat_vector<int32_t>({1, 2}),
        }),
        make_row_vector({
            make_flat_vector<int32_t>({1, 2}),
            make_flat_vector<int32_t>({1, 1}),
            make_flat_vector<int32_t>({1, 2}),
        }),
        make_row_vector({
            make_flat_vector<int32_t>({2, 2}),
            make_flat_vector<int32_t>({2, 3}),
            make_flat_vector<int32_t>({1, 2}),
        }),
    };

    createDuckDbTable(vectors);

    auto plan =
            PlanBuilder()
            .values(vectors)
            .markDistinct("c1_distinct", {"c0", "c1"})
            .markDistinct("c2_distinct", {"c0", "c2"})
            .singleAggregation(
                {"c0"}, {"sum(c1)", "sum(c2)"}, {"c1_distinct", "c2_distinct"})
            .planNode();

    AssertQueryBuilder(plan, duckDbQueryRunner_)
            .assertResults(
                "SELECT c0, sum(distinct c1), sum(distinct c2) FROM tmp GROUP BY 1");
}
